#include "net.hpp"

CustomDtaset::CustomDtaset(at::Tensor input_data, at::Tensor label_data)
{
    inputs_ = input_data;
    labels_ = label_data;
}

torch::data::Example<> CustomDtaset::get(size_t index)
{
    torch::Tensor Input = inputs_[index];
    torch::Tensor Label = labels_[index];
    return {Input.clone(), Label.clone()};
}

torch::optional<size_t> CustomDtaset::size() const { return labels_.size(0); }

Net::Net() {}

Net::Net(const at::Tensor& A, const at::Tensor& B, const at::Tensor& H, const int m,
         const int n)
{
    matA = register_parameter("matA", A, false);
    matB = register_parameter("matB", B, false);
    matH = register_parameter("matH", H, false);

    KalmanGain = register_parameter("KalmanGain", torch::ones({m, n}), true);
}

Net::~Net() {}

at::Tensor Net::forward(at::Tensor u, at::Tensor xtm1, at::Tensor z)
{
    // 预测
    torch::Tensor xt_ = torch::mm(u, matB) + torch::mm(xtm1, matA);

    // 校正
    torch::Tensor xt = torch::mm(z - torch::mm(xtm1, matH), KalmanGain) + xt_;

    return xt;
}
